import torch.utils.data as tud
import pandas as pd
from PIL import Image
from torchvision import transforms
import numpy as np

dataset_path = r'D:\datamining_pro\data/'  # type: str

CELESTIAL_BODY = {
    'qso': 0,
    'star': 1,
    'galaxy': 2
}


class CelestialDataset(tud.Dataset):
    def __init__(self, action='test', group_num=0):
        # We set the transform func like following set
        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Lambda(lambda x: x.repeat(3, 1, 1)), ])

        if action == 'train':
            if group_num == 0:
                # clean the header
                self.data = pd.read_csv(dataset_path + 'sets_v1_0.csv', header=None).iloc[1:20000, :2600].values
                self.labels = pd.read_csv(dataset_path + '训练集label.csv', header=None).iloc[1:20000, 1].values
            else:
                self.data = pd.read_csv(dataset_path + 'sets_v1_' + str(group_num) + '.csv',
                                        header=None).iloc[:20000, :2600].values
                self.labels = pd.read_csv(dataset_path + '训练集label.csv',
                                          header=None).iloc[(20000 * group_num):(20000 * group_num + 20000), 1].values
        elif action == 'test':
            data1 = pd.read_csv(dataset_path + 'sets_v1_7.csv', header=None).iloc[:20000, :2600].values
            data2 = pd.read_csv(dataset_path + 'sets_v1_8.csv', header=None).iloc[:20000, :2600].values
            self.data = np.vstack((data1, data2))
            self.labels = pd.read_csv(dataset_path + '验证集label.csv', header=None).iloc[1:40001, 1].values

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, index):
        single_label = self.labels[index]
        single_label = CELESTIAL_BODY[single_label]

        np_data = np.asarray(self.data[index]).reshape(26, 100).astype(float)
        data_max = np.max(np_data)
        data_min = np.min(np_data)
        np_data = (np_data - data_min) / ((data_max - data_min) + 0.0000001)
        single_data = Image.fromarray(np_data)
        if self.transform is not None:
            single_data = self.transform(single_data)

        return single_data, single_label
